#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Repmgr管理器模块
提供PostgreSQL复制管理功能
"""

import os
import re
import subprocess
from typing import Optional, List, Dict

from ..core.utils import (
    logger,
    run_command,
    detect_postgresql_version,
    is_windows,
    confirm_action,
    get_input_with_default
)
from ..core import config

class RepmgrManager:
    """PostgreSQL复制管理器类"""
    
    def __init__(self):
        self.pg_version = detect_postgresql_version()
        self.config_path = config.REPMGR_CONFIG_PATH
        self.bin_path = f"/usr/lib/postgresql/{self.pg_version}/bin" if self.pg_version else None
    
    def is_postgresql_running(self) -> bool:
        """检查PostgreSQL是否正在运行"""
        try:
            result = run_command(["pg_isready"], check=True, capture_output=True)
            return result.returncode == 0
        except Exception:
            logger.error("检查PostgreSQL状态时出错")
            return False
    
    def install_repmgr(self):
        """安装Repmgr复制管理工具"""
        version = get_input_with_default("输入要安装的repmgr版本（留空为最新版本）", "latest")
        
        try:
            if version == "latest":
                # 安装最新版本
                run_command(["sudo", "apt-get", "-y", "install", "repmgr"])
            else:
                # 安装指定版本
                run_command(["sudo", "apt-get", "-y", "install", f"repmgr={version}"])
                
            # 输出已安装的版本
            version_output = run_command(["repmgr", "--version"], capture_output=True)
            logger.info(f"repmgr安装成功。版本: {version_output.stdout.strip()}")
        except Exception as e:
            logger.error(f"安装repmgr失败: {e}")
    
    def create_repmgr_user_and_db(self):
        """创建repmgr超级用户和数据库"""
        if not self.pg_version:
            logger.error("无法检测PostgreSQL版本")
            return
            
        try:
            # 创建超级用户repmgr
            logger.info("正在创建超级用户'repmgr'...")
            run_command([
                "sudo", "-u", "postgres",
                f"/usr/lib/postgresql/{self.pg_version}/bin/createuser", "-s", "repmgr"
            ])
            
            # 创建专用数据库repmgr
            logger.info("正在创建'repmgr'数据库")
            run_command([
                "sudo", "-u", "postgres",
                f"/usr/lib/postgresql/{self.pg_version}/bin/createdb", "repmgr", "-O", "repmgr"
            ])
            
            # 设置repmgr用户的搜索路径
            logger.info("设置repmgr用户的搜索路径")
            run_command([
                "sudo", "-u", "postgres",
                f"/usr/lib/postgresql/{self.pg_version}/bin/psql", "-p", "5432", "-h", "127.0.0.1", "-U", "postgres",
                "-c", "ALTER USER repmgr SET search_path TO repmgr, \"$user\", public;"
            ])
            
            logger.info("超级用户'repmgr'和数据库'repmgr'已成功创建")
        except Exception as e:
            logger.error(f"创建用户和数据库失败: {e}")
    
    def create_repmgr_config(self):
        """创建Repmgr配置文件"""
        if not self.pg_version:
            logger.error("无法检测PostgreSQL版本")
            return
            
        config_content = f"""node_id=1
node_name=node1
conninfo='host=node1 user=repmgr dbname=repmgr connect_timeout=2'
data_directory='/var/lib/postgresql/{self.pg_version}/main'
pg_bindir='/usr/lib/postgresql/{self.pg_version}/bin/'
ssh_options='-q -o ConnectTimeout=10'
failover='manual' 
promote_command='/usr/lib/postgresql/{self.pg_version}/bin/repmgr standby promote -f {self.config_path} --log-to-file'
follow_command='/usr/lib/postgresql/{self.pg_version}/bin/repmgr standby follow -f {self.config_path} --log-to-file --upstream-node-id=%n'
service_start_command='sudo systemctl start postgresql@{self.pg_version}-main.service'
service_stop_command='sudo systemctl stop postgresql@{self.pg_version}-main.service'
service_restart_command='sudo systemctl restart postgresql@{self.pg_version}-main.service'
service_reload_command='sudo systemctl reload postgresql@{self.pg_version}-main.service'
repmgrd_service_start_command='sudo systemctl start repmgrd'
repmgrd_service_stop_command='sudo systemctl stop repmgrd'
service_promote_command='pg_ctlcluster {self.pg_version} main promote'
"""
        
        try:
            # 创建配置文件并写入内容
            config_dir = os.path.dirname(self.config_path)
            
            # 检查目录是否存在，如果不存在则创建
            if not os.path.exists(config_dir):
                os.makedirs(config_dir)
                
            with open(self.config_path, "w") as config_file:
                config_file.write(config_content)
                
            # 获取postgres用户的uid和gid（仅在非Windows系统）
            if not is_windows():
                try:
                    postgres_uid = int(subprocess.run(
                        ["id", "-u", "postgres"], 
                        capture_output=True, 
                        text=True).stdout.strip())
                    postgres_gid = int(subprocess.run(
                        ["id", "-g", "postgres"], 
                        capture_output=True, 
                        text=True).stdout.strip())
                    
                    # 设置文件的所有者为postgres用户
                    os.chown(self.config_path, postgres_uid, postgres_gid)
                    
                    # 设置文件权限为700
                    os.chmod(self.config_path, 0o700)
                except Exception as e:
                    logger.error(f"设置文件权限和所有者时出错: {e}")
            else:
                # Windows设置文件权限
                try:
                    os.chmod(self.config_path, 0o600)
                except Exception as e:
                    logger.error(f"设置文件权限时出错: {e}")
            
            logger.info(f"配置文件已成功创建: {self.config_path}")
        except Exception as e:
            logger.error(f"配置文件创建失败: {e}")
    
    def get_repmgr_binary(self) -> str:
        """获取repmgr命令的完整路径"""
        if self.bin_path:
            return os.path.join(self.bin_path, "repmgr")
        return "repmgr"
    
    def register_primary_node(self):
        """注册主节点"""
        if not self.pg_version:
            logger.error("无法检测PostgreSQL版本")
            return
            
        try:
            logger.info("正在将当前节点注册为主节点...")
            run_command([
                "sudo", "-u", "postgres",
                self.get_repmgr_binary(), "primary", "register",
                "-f", self.config_path,
                "--force", "--log-to-file"
            ])
            logger.info("当前节点已成功注册为主节点")
        except Exception as e:
            logger.error(f"注册主节点失败: {e}")
    
    def register_standby_node(self):
        """注册从节点"""
        if not self.pg_version:
            logger.error("无法检测PostgreSQL版本")
            return
            
        try:
            logger.info("正在将当前节点注册为从节点...")
            run_command([
                "sudo", "-u", "postgres",
                self.get_repmgr_binary(), "standby", "register",
                "-f", self.config_path,
                "--force"
            ])
            logger.info("当前节点已成功注册为从节点")
        except Exception as e:
            logger.error(f"注册从节点失败: {e}")
    
    def hand_switchover(self):
        """手动主备切换，将当前从节点升级为主节点，主节点降级为从节点"""
        if not self.pg_version:
            logger.error("无法检测PostgreSQL版本")
            return
            
        try:
            logger.info("正在将当前节点升级为主节点，主节点将降级为从节点...")
            run_command([
                "sudo", "-u", "postgres",
                self.get_repmgr_binary(), "standby", "switchover",
                "-f", self.config_path,
                "--siblings-follow"
            ])
            logger.info("主备切换完成，当前节点已是主节点")
        except Exception as e:
            logger.error(f"主备切换失败: {e}")
    
    def handle_old_primary_recovery(self):
        """裂脑处理，将旧主节点重新加入集群作为从节点"""
        if not self.pg_version:
            logger.error("无法检测PostgreSQL版本")
            return
            
        # 提示用户输入新主节点的IP地址
        new_primary_ip = get_input_with_default("请输入新主节点IP地址", "")
        
        if not new_primary_ip:
            logger.error("新主节点IP地址不能为空")
            return
            
        try:
            # 停止旧主节点服务
            logger.info("正在停止当前节点数据库...")
            old_primary_service = f"postgresql@{self.pg_version}-main.service"
            run_command(["sudo", "systemctl", "stop", old_primary_service])
            
            # 以备节点身份重新加入集群
            logger.info("正在将当前节点以从库身份加入集群...")
            run_command([
                "sudo", "-u", "postgres",
                self.get_repmgr_binary(), "-f", self.config_path,
                "node", "rejoin",
                "-d", f"host={new_primary_ip} dbname=repmgr user=repmgr",
                "--force-rewind"
            ])
            
            logger.info("当前节点已作为从节点重新加入集群")
        except Exception as e:
            logger.error(f"裂脑处理失败: {e}")
    
    def show_cluster(self):
        """查看集群状态"""
        if not self.pg_version:
            logger.error("无法检测PostgreSQL版本")
            return
            
        try:
            run_command([
                "sudo", "-u", "postgres",
                self.get_repmgr_binary(), "-f", self.config_path,
                "cluster", "show"
            ])
        except Exception as e:
            logger.error(f"查看集群信息失败: {e}")
    
    def clone_standby(self):
        """从主节点克隆数据到本地，创建从节点"""
        if not self.pg_version:
            logger.error("无法检测PostgreSQL版本")
            return
            
        primary_host = get_input_with_default("请输入主节点地址", "")
        
        if not primary_host:
            logger.error("主节点地址不能为空")
            return
            
        # 准备克隆命令
        standby_clone_command = [
            "sudo", "-u", "postgres",
            self.get_repmgr_binary(), "-f", self.config_path,
            "standby", "clone", "-h", primary_host, 
            "-U", config.DEFAULT_REPMGR_USER, 
            "-d", config.DEFAULT_REPMGR_DATABASE
        ]
        
        # 验证克隆条件并执行克隆
        try:
            logger.info("正在验证是否满足克隆条件...")
            run_command(standby_clone_command + ["--dry-run"])
            
            if confirm_action("根据输出结果，是否继续克隆操作?", True):
                run_command(standby_clone_command)
                logger.info("克隆成功完成")
            else:
                logger.info("用户取消了克隆操作")
        except Exception as e:
            logger.error(f"克隆操作失败: {e}")
    
    def create_sudoers_file(self):
        """为postgres用户添加必要的sudo权限"""
        if not self.pg_version:
            logger.error("无法检测PostgreSQL版本")
            return
            
        config_content = f"""
postgres ALL = (root) NOPASSWD:/usr/bin/systemctl start postgresql@{self.pg_version}-main.service,/usr/bin/systemctl restart postgresql@{self.pg_version}-main.service,/usr/bin/systemctl stop postgresql@{self.pg_version}-main.service,/usr/bin/systemctl reload postgresql@{self.pg_version}-main.service,/usr/bin/systemctl status postgresql@{self.pg_version}-main.service,/usr/bin/systemctl start repmgrd,/usr/bin/systemctl stop repmgrd,/usr/bin/systemctl reload repmgrd,/usr/bin/systemctl status repmgrd,/usr/bin/systemctl restart repmgrd
"""
        
        try:
            # 将配置内容写入文件
            with open("/etc/sudoers.d/postgres", "w") as file:
                file.write(config_content)
                
            # 检查sudoers文件语法
            result = run_command(["sudo", "visudo", "-c", "-f", "/etc/sudoers.d/postgres"], capture_output=True)
            if result.returncode != 0:
                logger.error("sudoers文件存在语法错误，请检查")
                return False
                
            # 设置文件权限
            run_command(["sudo", "chmod", "440", "/etc/sudoers.d/postgres"])
            run_command(["sudo", "chown", "root:root", "/etc/sudoers.d/postgres"])
            
            logger.info("sudoers文件创建并设置权限成功")
            return True
        except Exception as e:
            logger.error(f"创建sudoers文件时出错: {e}")
            return False
    
    def repmgr_nodestatus(self):
        """检查当前节点健康状态"""
        try:
            run_command([
                "sudo", "-u", "postgres",
                "repmgr", "-f", self.config_path, "node", "check"
            ])
        except Exception as e:
            logger.error(f"健康状态检查失败: {e}")
    
    def standby_follow(self):
        """手动将当前从节点跟随新主节点"""
        if not self.is_postgresql_running():
            logger.error("本机PostgreSQL未运行，无法执行操作")
            return
            
        upstream_node_id = get_input_with_default("请输入要连接的上游节点ID（如果没有，请直接回车跳过）", "")
        
        command = [
            "sudo", "-u", "postgres",
            self.get_repmgr_binary(), "-f", self.config_path,
            "standby", "follow"
        ]
        
        if upstream_node_id:
            command.extend(["--upstream-node-id", upstream_node_id])
            
        try:
            logger.info("正在设置当前节点跟随新的主节点...")
            run_command(command)
            logger.info("设置成功！当前节点已跟随新的主节点")
        except Exception as e:
            logger.error(f"设置失败: {e}")
    
    def standby_promote(self):
        """将当前从节点升级为主节点（主节点故障时使用）"""
        if not self.is_postgresql_running():
            logger.error("本机PostgreSQL未运行，无法执行操作")
            return
            
        command = [
            "sudo", "-u", "postgres",
            self.get_repmgr_binary(), "-f", self.config_path,
            "standby", "promote", "--siblings-follow"
        ]
        
        try:
            logger.info("正在升级当前节点为主节点...")
            run_command(command)
            logger.info("升级成功！当前节点为主节点")
        except Exception as e:
            logger.error(f"升级失败: {e}")
    
    def create_repmgrd_service(self):
        """创建并启动repmgrd服务"""
        if not self.pg_version:
            logger.error("无法检测PostgreSQL版本")
            return
            
        # 删除旧的init.d文件
        if os.path.exists("/etc/init.d/repmgrd"):
            try:
                run_command(["sudo", "rm", "-f", "/etc/init.d/repmgrd"])
            except Exception as e:
                logger.error(f"删除旧服务文件时出错: {e}")
                return
                
        # 生成systemd服务文件内容
        service_content = f"""\
[Unit]
Description=A replication manager and failover management tool for PostgreSQL
After=syslog.target
After=network.target
After=postgresql.service

[Service]
Type=forking
User=postgres
Group=postgres

# Use the /run/postgresql directory for the PID file
ExecStartPre=/bin/mkdir -p /run/postgresql
ExecStartPre=/bin/chown postgres:postgres /run/postgresql

# PID file
PIDFile=/run/postgresql/repmgrd.pid

# Location of repmgr config file
Environment=REPMGRDCONF={self.config_path}

# Start repmgrd with the correct configuration and PID file
ExecStart=/usr/lib/postgresql/{self.pg_version}/bin/repmgrd -f ${{REPMGRDCONF}} --daemonize=true --pid-file=/run/postgresql/repmgrd.pid --verbose

# Commands to stop and reload repmgrd
ExecStop=/usr/bin/kill -TERM $MAINPID
ExecReload=/usr/bin/kill -HUP $MAINPID

# Timeout for startup and shutdown
TimeoutSec=300

[Install]
WantedBy=multi-user.target
"""
        
        # 写入服务文件
        try:
            with open("/etc/systemd/system/repmgrd.service", "w") as f:
                f.write(service_content)
            logger.info("repmgrd.service文件创建成功")
        except Exception as e:
            logger.error(f"创建repmgrd.service文件时出错: {e}")
            return
            
        # 重新加载systemd配置
        try:
            run_command(["sudo", "systemctl", "daemon-reload"])
            run_command(["sudo", "systemctl", "enable", "repmgrd.service"])
            run_command(["sudo", "systemctl", "start", "repmgrd.service"])
            logger.info("repmgrd服务已启用并启动")
        except Exception as e:
            logger.error(f"启动repmgrd服务时出错: {e}")
    
    def repair_repmgr_extension(self):
        """修复repmgr扩展"""
        try:
            logger.info("正在修复repmgr扩展...")
            
            # 创建扩展
            run_command([
                "sudo", "-u", "postgres", "psql", "-d", config.DEFAULT_REPMGR_DATABASE,
                "-c", "CREATE EXTENSION repmgr;"
            ])
            
            # 查询扩展是否成功安装
            result = run_command([
                "sudo", "-u", "postgres", "psql", "-d", config.DEFAULT_REPMGR_DATABASE,
                "-c", "SELECT * FROM pg_extension WHERE extname = 'repmgr';"
            ], capture_output=True)
            
            logger.info(f"查询扩展结果:\n{result.stdout}")
            
            # 检查输出中是否包含'repmgr'
            if "repmgr" in result.stdout:
                logger.info("repmgr扩展安装成功")
            else:
                logger.error("repmgr扩展安装失败或未找到")
        except Exception as e:
            logger.error(f"修复repmgr扩展时出错: {e}")
    
    def regist_witness_node(self):
        """注册见证节点"""
        primary_ip = get_input_with_default("请输入主节点IP地址", "")
        
        if not primary_ip:
            logger.error("主节点IP地址不能为空")
            return
            
        try:
            run_command([
                "sudo", "-u", "postgres",
                "repmgr", "-f", self.config_path, "witness", "register",
                "-h", primary_ip, "-F"
            ])
            logger.info("见证节点注册成功")
        except Exception as e:
            logger.error(f"注册见证节点失败: {e}") 